Repository navigation
Add channel_axis to rotate and affine_transform - #142
krakken190 wants to merge 5 commits into
Conversation
krakken190
changed the title
|
Hi @claudiofantacci , I've resolved the merge conflict with master. Ready for review/merge whenever you get a chance! |
|
@claudiofantacci Sorry, my merge-conflict resolution in the GitHub UI dropped the |
| @parameterized.product( | ||
| mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), | ||
| ("constant", 1)], | ||
| channel_axis=(0, -3), |
There was a problem hiding this comment.
The use of -3 seems a bit unnecessary here. Can make it explicit?
There was a problem hiding this comment.
Done, the tests now use channel_axis=(-1, 0), i.e. HWC and CHW.
| ("constant", 1)], | ||
| channel_axis=(0, -3), | ||
| ) | ||
| def test_affine_transform_channel_axis(self, mode_and_order, channel_axis): |
There was a problem hiding this comment.
Can we add also the different ways of passing the dimensions/offset? See how test_affine_transform is done, and we just add channel axis to the product of the parameterized test.
There was a problem hiding this comment.
Done. channel_axis is now part of the parameterized product in test_affine_transform, so all matrix/offset forms are checked against SciPy for both HWC and CHW. I removed the separate test.
| mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), | ||
| ("constant", 1)], | ||
| angle=(np.pi / 2, -np.pi / 4), | ||
| channel_axis=(0, -3), |
| angle=(np.pi / 2, -np.pi / 4), | ||
| channel_axis=(0, -3), | ||
| ) | ||
| def test_rotate_channel_axis(self, mode_and_order, angle, channel_axis): |
There was a problem hiding this comment.
Can we test it in the same manner we do test_rotate, we just need to add channel axis to that test probably. Scipy as reference function has axes of which oens you want to rotate.
There was a problem hiding this comment.
agree, now test_rotate takes channel_axis and uses SciPy's axes argument as the reference: (1, 0) for HWC and (2, 1) for CHW. I removed the separate test.
|
@pabloduque0 Thanks for the review! All comments are addressed in 7dab6fa: channel_axis is now tested inside the existing parameterized test_affine_transform and test_rotate, and the separate tests are removed. I also merged the latest master; all tests pass locally (1943 passed). Could you approve the workflows when you have a moment? |
Fixes #141.
Problem
rotateandaffine_transformsay they support HWC or CHW images, but they always treat the first two array axes as height and width. For CHW inputs this rotates or shifts across the channel axis and silently returns a corrupted image.Fix
channel_axis: int = -1toaffine_transform. When it is not-1, the channel axis is moved to the end before the transform and moved back afterwards. This is the same patternelastic_deformationuses.matrixandoffsetare always interpreted in HWC order, independent ofchannel_axis, so one transform can be applied to either layout. The docstring now says this.channel_axis: int = -1torotate. It computes the rotation centre from the HWC-ordered shape and forwardschannel_axistoaffine_transform.-1keeps the existing behaviour, so current callers are unaffected.Tests
Added to
TestCustom:test_affine_transform_channel_axis: fornearest/0,nearest/1,mirror/1 andconstant/1, withchannel_axis0 and -3, the CHW result equals the HWC result moved to channels-first.test_rotate_channel_axis: the same check for two angles, run underjax.jit.Validation
pytest dm_pix/_src/augment_test.py: 886 passedpytest dm_pix/_src/api_test.py: 80 passedchannel_axis=0, CHW output matches HWC exactly (max diff 0.0) for both functions. Without it, the difference was 0.99997 before and after the change.Tested on CPU with Python 3.12.2, JAX/jaxlib 0.12.0, NumPy 2.5.4, SciPy 1.18.1 and TensorFlow 2.21.0. GPU/TPU not tested.